package com.tianji.aigc.controller;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.Arrays;
import java.util.List;

/**
 * @version: 1.0
 * @Author: ljy
 * @description: 向量数据库Controller
 * @date: 2025-09-03 14:58
 */

@RestController
@RequestMapping("/embedding")
@Slf4j
public class EmbeddingController {

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private EmbeddingModel embeddingModel;

    /**
     * 文本转向量
     * @param message
     * @return
     */
    @GetMapping
    public EmbeddingResponse textToEmbedding(String message){
        EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(Arrays.asList(message));
        return embeddingResponse;
    }


    /**
     * 文本保存到向量库
     *
     * @param messageList
     */
    @PostMapping
    public void saveVectorStore(@RequestParam("messages") List<String> messageList) {
        //字符串消息转为文档消息
        List<Document> documentList = messageList.stream().map(
                m -> Document.builder().text(m).build()
        ).toList();

        //添加文档消息到向量库
        vectorStore.add(documentList);

        log.info("[向量库添加数据]添加完毕，添加成功{}条",messageList.size());
    }

    @DeleteMapping
    public void deleteVectorStore(@RequestParam("ids") List<String> ids) {
        // 删除向量数据库中的数据
        this.vectorStore.delete(ids);
    }

    @GetMapping("/search")
    public List<Document> search(@RequestParam("message") String message) {
        return this.vectorStore.similaritySearch(SearchRequest.builder().query(message).topK(5).build());
    }

    @GetMapping("/search/all")
    public List<Document> searchAll() {
        // 搜索全部数据
        return this.vectorStore.similaritySearch(SearchRequest.builder().query("").topK(999).build());
    }
}
